import cv2
import numpy as np

# 加载图像
image = cv2.imread('piano512.PNG', cv2.IMREAD_GRAYSCALE)
watermark = cv2.imread('lsc_sign512.jpg', cv2.IMREAD_GRAYSCALE)

# 执行DFT变换
image_dft = np.fft.fft2(image)
watermark_dft = np.fft.fft2(watermark)

# 将水印嵌入到频域图像中
alpha = 0.1  # 调整水印强度
watermarked_dft = image_dft + alpha * watermark_dft

# 执行逆DFT变换
watermarked_image = np.fft.ifft2(watermarked_dft).real

# 显示原始图像、水印和水印后的图像
cv2.imshow('Original Image', image)
cv2.imshow('Watermark', watermark)
cv2.imshow('Watermarked Image', watermarked_image.astype(np.uint8))
cv2.waitKey(0)
cv2.destroyAllWindows()

# 提取水印
extracted_dft = (watermarked_dft - image_dft) / alpha
extracted_image = np.fft.ifft2(extracted_dft).real

# 显示提取的水印
cv2.imshow('Extracted Watermark', extracted_image.astype(np.uint8))
cv2.waitKey(0)
cv2.destroyAllWindows()
